import math
import numpy as np

if '__main__' == __name__:
    np.set_printoptions(threshold=np.inf)

    scale = 2.0
    N = 8
    t = np.linspace(-scale, scale, N)
    chunk = 1024 * 64
    n_iters = int(math.ceil(N * N * N / chunk))

    query_pts = np.stack(np.meshgrid(t, t, t), -1).astype(np.float32)
    print('points shape:', query_pts.shape)
    print(query_pts)
    sh = query_pts.shape
    flat = query_pts.reshape([-1, 3])
    print('flat shape:', flat.shape)

    def calc_fn(points):
        points -= 1.
        sigma = points.max(axis=-1)
        return sigma

    raw_list = []
    for i in range(n_iters):
        points = flat[i * chunk:(i + 1) * chunk]
        raw_entry = np.array(calc_fn(points))
        raw_list.append(raw_entry)
    raw = np.concatenate(raw_list, axis=0)
    print('raw.shape', raw.shape)
